import numpy as np
import tensorflow as tf
import dlib
# import matplotlib.pyplot as plt
import cv2 as cv
import os
import pickle
# from scipy import misc
import imageio
from queue import Queue
from multiprocessing import Queue as process_queue
from multiprocessing import Process
import time

tf.compat.v1.disable_eager_execution()

# # GPU设置
# tf.debugging.set_log_device_placement(True)
# gpus = tf.config.experimental.list_physical_devices('GPU')
# # 设置第二块可见
# tf.config.experimental.set_visible_devices(gpus[1], 'GPU')
# # 设置内存自增长
# for gpu in gpus:
#     tf.config.experimental.set_memory_growth(gpu, True)

# # 可能 GPU 的内存会不够，因此加上这句
# from tensorflow.compat.v1 import ConfigProto
# from tensorflow.compat.v1 import InteractiveSession
# from tensorflow.com
# config = ConfigProto()
# config.gpu_options.allow_growth = True
# session = InteractiveSession(config=config)
tf.keras.backend.clear_session()

class FaceNet:
    '''
        人脸识别网络实现类
        Args:
            self.queue:阻塞队列，有人注册他的图片就会被放进队列中，然后依次从队列中提取特征
            self.img_size:模型输入的图片需要是160x160x3的
            embeddings_txt_path: 注册图片特征保存的位置（看后端想放在哪里，只需要传文件夹路径进来就可以了）
            worker_id_path: 与注册图片特征保存相对应的工人id保存文件路径
            models_base_path: 模型存放的路径
            threshold:阈值
            self.inference:模型推理
    '''
    def __init__(self, worker_id_dir="E:\\zhangshiWorkSpace\\westudy\\facenet-tf2\\test", \
                 embeddings_txt_dir="E:\\zhangshiWorkSpace\\westudy\\facenet-tf2\\test", \
                 models_base_path="E:\\zhangshiWorkSpace\\westudy\\facenet-tf2\\src\\models", \
                 threshold=0.33):
        self.queue = Queue()
        self.img_size=160
        self.embeddings_txt_path = os.path.join(embeddings_txt_dir, "register_embeddings.txt")
        self.worker_id_path = os.path.join(worker_id_dir, "worker_ids.txt")
        self.threshold = threshold
        self.models_base_path = models_base_path
        model = tf.saved_model.load(self.models_base_path)
        self.inference = model.signatures['serving_default']
        # 如果文件不存在，那么创建
        self.file_is_exist(embeddings_txt_dir, worker_id_dir)
        # 因为这两个文件经常用到，需不需要就是每次实例化就展示出来效率会比较高
#         self.register_embeddings = np.loadtxt(embeddings_txt_path)
#         self.worker_id = np.loadtxt(worker_id_path)

    def register(self, worker_id, image_paths):
        # 每次有人注册，可能会同时有很多人注册，所以先放进队列里
        self.queue.put((worker_id, image_paths))
        self.register_()

    def register_(self):
        # queue 是一个阻塞函数，一个一个拿出来
        worker_id, image_path = self.queue.get()
        # 通过路径把图片路径加载进来，并且对人脸框进行对其和裁剪
        image_data_arr = self.load_and_align_data(image_path)
        # print('图片:', image_data_arr)
        # 使用模型对数据进行推理
        register_embeddings = self.model_infer(image_data_arr)
        # 得到的特征保存到文件中去
        with open(self.embeddings_txt_path,"ab+") as f:
            np.savetxt(f, register_embeddings, fmt='%.6f')
        with open(self.worker_id_path,"ab+") as f:
            for i in range(register_embeddings.shape[0]):
                np.savetxt(f, np.array([worker_id]), fmt="%i")
        print("{}张人脸注册成功".format(image_data_arr.shape[0]))

    def distance(self, login_embeddings, register_embeddings, q, worker_ids):
        # 使用欧几里得距离计算公式
        diff = np.subtract(login_embeddings, register_embeddings)
        dist = np.sum(np.square(diff), 1)
        # 通过下标获取当前登录人的 worker_id
        worker_id = worker_ids[np.less(dist, self.threshold)]
        if q:
            q.put(worker_id)
        return worker_id

    def login(self, image_paths):
        # 存放多进程结果
        q = process_queue()
        start_time = time.time()
        image_data_arr = self.load_and_align_data(image_paths)
        end_time = time.time()
        print("加载图片和裁剪图片消耗时间为{:.4f}".format(end_time-start_time))
        start_time = time.time()
        login_embeddings = self.model_infer(image_data_arr)
        end_time = time.time()
        print("模型推理消耗消耗时间为{:.4f}".format(end_time-start_time))
        # 加载已经注册过的人脸文件
        register_embeddings = np.loadtxt(self.embeddings_txt_path)
        worker_ids = np.loadtxt(self.worker_id_path)

        print(register_embeddings.shape)
        print(worker_ids.shape)
#         workers = []
        # 如果登录的是多个人，那么开启多线程的距离计算，加快速度
        start_time = time.time()
        for i in range(login_embeddings.shape[0]):
            _ = self.distance(login_embeddings[i], register_embeddings, q, worker_ids)
#             t = Process(target=self.distance, args=(login_embeddings[i], register_embeddings, q, worker_ids))
#             t.start()
#             workers.append(t)
        # 返回多进程处理完的结果
        end_time = time.time()
        print("距离计算消耗消耗时间为{:.4f}".format(end_time-start_time))
        res = []
        for _ in range(login_embeddings.shape[0]):
            res.extend(q.get())
        return res

    def model_infer(self, img_data_arr):
        img_data_arr = tf.cast(img_data_arr, tf.float32)
        # 归一化
        img_data_arr /= 255.0
        embeddings = self.inference(img_data_arr)
        embeddings = embeddings['output_0']
        # 进行正则化
        embeddings = tf.nn.l2_normalize(embeddings, 1, 1e-10, name='embeddings')

        print("模型推理结束")

        # 把 Tensor 转化为 numpy，方便后续计算
        sess = tf.compat.v1.InteractiveSession()
        print("1111111111111111111")
        sess.run(tf.compat.v1.global_variables_initializer())
        print("22222222222222222222222")
        embeddings = sess.run(embeddings)
        return embeddings

    def file_is_exist(self, embeddings_txt_dir, worker_id_dir):
        if not os.path.exists(embeddings_txt_dir):
            os.mkdir(embeddings_txt_dir)
        if not os.path.exists(worker_id_dir):
            os.mkdir(worker_id_dir)

    def load_and_align_data(self, image_path):
        '''
        因为我们的网络输入的 shape 需要是 160x160 的
        所以，每次新图片进来，我们都需要进行裁剪
        Args:
            输入进来的是图片的路径,可能是列表或者单个文件
        Return：
            返回图片的数据(可能是一张图片有多个人脸)，所以是个数组
        '''
        # 使用 dlib 库进行人脸对齐,首先初始化人脸探测器
        detector = dlib.get_frontal_face_detector()
        img_list = []
        # 把每张图片一次拿出来
        for image in image_path:
            img_data = imageio.imread(image, pilmode='RGB')
            # 进行人脸框检测
            dets = detector(img_data, 1)
            print('人数：', len(dets))
            # 如果检测不到人脸，那么跳过
            if dets.__len__() == 0:
                continue
            # 一张图片可能有多张人脸
            for i in range(len(dets)):
                face = dets[i]
                x1 = face.left()
                y1 = face.top()
                x2 = face.right()
                y2 = face.bottom()
                # 人脸框可能会出现比较小，为了拿到整张人脸，我们进行适当缩放
                y1 = int(y1 - (y2 - y1) * 0.3)
                x1 = int(x1 - (x2 - x1) * 0.05)
                x2 = int(x2 + (x2 - x1) * 0.05)
                y2 = y2
                # 裁剪图片
                im_crop_data = img_data[y1:y2, x1:x2]
                # 裁剪成 160 x 160 的格式
                im_crop_data = cv.resize(im_crop_data, (self.img_size, self.img_size))
                img_list.append(im_crop_data)
        # 把两张图片返回回去
        return np.array(img_list)



# if __name__ == '__main__':
#     # 实例化facenet类
#     facenet = facenet()
#     # image_paths 传进来的必须是 list 格式的,因为注册的时候人脸是多张的
#     # worker_id 是工人的主键 id
#     # 产生测试数据集
# #     for i in range(659, 10000, 3):
# #         facenet.register(worker_id=i, image_paths=["E:\\zhangshiWorkSpace\\westudy\\facenet-tf2\\test\\img_1.png"])
# #         facenet.register(worker_id=i+1, image_paths=["E:\\zhangshiWorkSpace\\westudy\\facenet-tf2\\test\\img_2.png"])
# #         facenet.register(worker_id=i+2, image_paths=["E:\\zhangshiWorkSpace\\westudy\\facenet-tf2\\test\\img_3.png"])
#     start_time = time.time()
#     # 识别的时候，只需要传进图片路径，需要是list；返回worker_id,可能是多个
#     worker_id_list = facenet.login(image_paths=["E:\\zhangshiWorkSpace\\westudy\\facenet-tf2\\test\\img_3.png"])
#     print(worker_id_list)
#     end_time = time.time()
#     print("消耗时间为{:.4f}".format(end_time-start_time))

# print('模型初始化')
# model = FaceNet(
#     worker_id_dir='D:/WeChat/download/WeChat Files/wxid_hyxc5wzmnklm12/FileStorage/File/2022-01/test',
#     embeddings_txt_dir='D:/WeChat/download/WeChat Files/wxid_hyxc5wzmnklm12/FileStorage/File/2022-01/test',
#     models_base_path='D:\\Code\\Python_code\\test\\myweb\\myapp\\models')

# print('\n新员工注册')
# model.register(worker_id=1, image_paths=['media\icon\8014f4bb0d93931.png'])

# print('\n员工识别')
# id_lst = model.login(image_paths=['media\icon\8014f4bb0d93931.png'])
# print('员工id:', int(id_lst[0]))